import os
from PIL import Image
import torch
from torch.utils.data import Dataset

class CustomImageFolder(Dataset):
    def __init__(self, root, transform=None):
        self.root = root
        self.transform = transform
        # self.classes = sorted([d.name for d in os.scandir(root) if d.is_dir()])
        self.classes = ['drink', 'listen', 'phone', 'trance', 'write']
        self.class_to_idx = {'drink': 0, 'listen': 1, 'phone': 2, 'trance': 3, 'write': 4}
        self.samples = []

        for target_class in sorted(self.class_to_idx.keys()):
            class_index = self.class_to_idx[target_class]
            target_dir = os.path.join(self.root, target_class)
            if not os.path.isdir(target_dir):
                continue

            for vid_name in sorted(os.listdir(target_dir)):
                vid_dir = os.path.join(target_dir, vid_name)
                if not os.path.isdir(vid_dir):
                    continue

                img_paths = sorted([os.path.join(vid_dir, f) for f in os.listdir(vid_dir) if f.endswith('.jpg')])
                # 将一个视频文件夹中的所有图片都放在一个数组中，对应一个标签
                video_data = []
                for img_path in img_paths:
                    img = Image.open(img_path).convert('RGB')
                    if self.transform is not None:
                        img = self.transform(img)
                    video_data.append(img)

                # 对视频数据进行扩充
                while len(video_data) < 40:
                    video_data.append(video_data[-1])

                while len(video_data) != 40:
                    del video_data[-1]

                self.samples.append((video_data, class_index))

        self.targets = [s[1] for s in self.samples]

    def __getitem__(self, index):
        video_data, target = self.samples[index]
        return torch.stack(video_data), target

    def __len__(self):
        return len(self.samples)


class TestDataset(Dataset):
    def __init__(self, root, transform=None):
        self.root = root
        self.transform = transform
        self.samples = []
        self.classes = ['drink', 'listen', 'phone', 'trance', 'write']

        # vid_dir = os.path.join(root, str(index))
        vid_dir = root
        if not os.path.isdir(vid_dir):
            raise FileNotFoundError(f"Directory '{vid_dir}' not found.")

        img_paths = sorted([os.path.join(vid_dir, f) for f in os.listdir(vid_dir) if f.endswith('.jpg')])

        # 将一个视频文件夹中的所有图片都放在一个数组中
        video_data = []
        for img_path in img_paths:
            img = Image.open(img_path).convert('RGB')
            if self.transform is not None:
                img = self.transform(img)
            video_data.append(img)

        # 对视频数据进行扩充
        while len(video_data) < 40:
            video_data.append(video_data[-1])

        while len(video_data) != 40:
            del video_data[-1]

        self.samples.append((video_data, 0))

        # self.targets = [s[1] for s in self.samples]

    def __getitem__(self, index):
        video_data, target = self.samples[index]
        return torch.stack(video_data), target

    def __len__(self):
        return len(self.samples)

